/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
package com.webank.tikv.operation;

import com.webank.tikv.exception.GrpcException;
import com.webank.tikv.exception.RegionException;
import com.webank.tikv.exception.TiKVException;
import com.webank.tikv.kvproto.Errorpb;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.region.RegionErrorReceiver;
import com.webank.tikv.region.RegionManager;
import com.webank.tikv.region.TiRegion;
import com.webank.tikv.region.TiStore;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.KeyUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

import static com.webank.tikv.util.BackOffer.BoType.BoDiskFull;
import static com.webank.tikv.util.BackOffer.BoType.BoMaxRegionNotInitialized;
import static com.webank.tikv.util.BackOffer.BoType.BoMaxTsNotSynced;
import static com.webank.tikv.util.BackOffer.BoType.BoRegionMiss;
import static com.webank.tikv.util.BackOffer.BoType.BoRegionRecoveryInProgress;
import static com.webank.tikv.util.BackOffer.BoType.BoServerBusy;
import static com.webank.tikv.util.BackOffer.BoType.BoStaleCmd;
import static com.webank.tikv.util.BackOffer.BoType.BoTiKVRPC;
import static com.webank.tikv.util.BackOffer.BoType.BoTsoBatchUsedUp;
import static com.webank.tikv.util.BackOffer.BoType.BoUpdateLeader;

public class RegionErrorHandler<RespT> implements ErrorHandler<RespT> {

    private static final Logger LOGGER = LogManager.getLogger(RegionErrorHandler.class);

    // if a store does not have leader currently, store id is set to 0
    public static final int NO_LEADER_STORE_ID = 0;
    private final Function<RespT, Errorpb.Error> regionErrorFunc;
    private final RegionManager regionManager;
    private final RegionErrorReceiver recv;

    public RegionErrorHandler(RegionManager regionManager, RegionErrorReceiver recv,
                              Function<RespT, Errorpb.Error> regionErrorFunc) {
        this.recv = recv;
        this.regionManager = regionManager;
        this.regionErrorFunc = regionErrorFunc;
    }

    @Override
    public boolean handleResponseError(BackOffer backOffer, RespT resp) {
        if (resp == null) {
            String msg = String.format("Request Failed with unknown reason for [%s]", recv.getRegion());
            return handleRequestError(backOffer, new GrpcException(msg));
        }
        Errorpb.Error error = null;
        if (regionErrorFunc != null) {
            error = regionErrorFunc.apply(resp);
        }
        if (error == null) {
            return false;
        }
        long sleep = handleRegionError(backOffer, error);
        if (sleep > 0) {
            try {
                Thread.sleep(sleep);
            } catch (InterruptedException ignored) {
            }
            return true;
        }
        return false;
    }

    @Override
    public long calNextRetryTime(BackOffer backOffer, RespT resp) {
        if (resp == null) {
            String msg = String.format("Request Failed with unknown reason for [%s]", recv.getRegion());
            if (handleRequestError(backOffer, new GrpcException(msg))) {
                return 0;
            }
        }
        Errorpb.Error error = null;
        if (regionErrorFunc != null) {
            error = regionErrorFunc.apply(resp);
        }
        if (error != null) {
            return handleRegionError(backOffer, error);
        }
        return -1;
    }

    private long handleRegionError(BackOffer backOffer, Errorpb.Error error) {
        TiRegion region = recv.getRegion();
        if (error.hasNotLeader()) {
            // this error is reported from raft store:
            // peer of current request is not leader, the following might be its causes:
            // 1. cache is outdated, region has changed its leader, can be solved by re-fetching from PD
            // 2. leader of current region is missing, need to wait and then fetch region info from PD
            long newStoreId = error.getNotLeader().getLeader().getStoreId();
            LOGGER.warn("tikv NotLeader with region:{}, new store:{}", region, newStoreId);

            if (newStoreId != NO_LEADER_STORE_ID) {
                // If update leader fails, we need to fetch new region info from pd,
                // and re-split key range for new region. Setting retry to false will
                // stop retry and enter handleCopResponse logic, which would use RegionMiss
                // backOff strategy to wait, fetch new region and re-split key range.
                // onNotLeader is only needed when updateLeader succeeds, thus switch
                // to a new store address.
                TiRegion newRegion = regionManager.updateLeader(region, newStoreId);
                if (newRegion != null && recv.onNotLeader(newRegion, backOffer)) {
                    regionManager.invalidRegion(region);
                }
                return backOffer.nextBackOffTime(BoUpdateLeader, new GrpcException(error.toString()));
            }

            // if there's current no leader, we do not trigger update pd cache logic
            // since issuing store = NO_LEADER_STORE_ID requests to pd will definitely fail.
            LOGGER.warn("tikv received zero store id with region:{}", region);
            return -1;
        } else if (error.hasStoreNotMatch()) {
            // this error is reported from raft store:
            // store_id requested at the moment is inconsistent with that expected
            // Solution：re-fetch from PD
            long actualStoreId = error.getStoreNotMatch().getActualStoreId();
            TiStore store = recv.getStore();
            LOGGER.warn("tikv StoreNotMatch with region:{}, store:{}, actual:{}",
                region.getId(), store, actualStoreId);
            regionManager.invalidRegion(region);
            regionManager.invalidStore(store.getId());
            return backOffer.nextBackOffTime(BoRegionMiss, new TiKVException(error.getMessage()));
        } else if (error.hasEpochNotMatch()) {
            LOGGER.warn("tikv EpochNotMatch with region:{}, store:{}", region, recv.getStore());
            return onRegionEpochNotMatch(backOffer, error);
        } else if (error.hasServerIsBusy()) {
            // this error is reported from kv: will occur when write pressure is high
            LOGGER.warn("tikv ServerBusy with region:{}", region);
            return backOffer.nextBackOffTime(BoServerBusy, new TiKVException(error.getMessage()));
        } else if (error.hasStaleCommand()) {
            // this error is reported from raft store: command outdated
            LOGGER.warn("tikv StaleCommand with region:{}", region);
            return backOffer.nextBackOffTime(BoStaleCmd, new TiKVException(error.getMessage()));
        } else if (error.hasKeyNotInRegion()) {
            Errorpb.KeyNotInRegion notInRegion = error.getKeyNotInRegion();
            LOGGER.warn("tikv KeyNotInRegion:[{}, {}], invalid region:{}",
                notInRegion.getRegionId(), KeyUtils.bytesToHex(notInRegion.getKey().toByteArray()), region);
            regionManager.invalidRegion(region);
            return backOffer.nextBackOffTime(BoRegionMiss, new TiKVException(error.getMessage()));
        } else if (error.hasRegionNotFound()) {
            Errorpb.RegionNotFound regionNotFound = error.getRegionNotFound();
            LOGGER.warn("tikv RegionNotFound with region:{} on:{}", regionNotFound.getRegionId(), region);
            regionManager.invalidRegion(region);
            return backOffer.nextBackOffTime(BoRegionMiss, new TiKVException(error.getMessage()));
        } else if (error.hasRegionNotInitialized()) {
            LOGGER.warn("tikv RegionNotInitialized with region:{}", error.getRegionNotInitialized().getRegionId());
            return backOffer.nextBackOffTime(BoMaxRegionNotInitialized, new TiKVException(error.getMessage()));
        } else if (error.getMessage().contains("TSO batch")) {
            // The tso cache is used up in TiKV servers, we should backoff and wait its cache is renewed.
            LOGGER.warn("tikv TSO Batch used up with region:{}", region);
            return backOffer.nextBackOffTime(BoTsoBatchUsedUp, new GrpcException(error.getMessage()));
        } else if (error.getMessage().contains("Raft raft: proposal dropped")) {
            // retry if raft proposal is dropped, it indicates the store is in the middle of transition
            return backOffer.nextBackOffTime(BoTiKVRPC, new TiKVException(error.getMessage()));
        } else if (error.hasRaftEntryTooLarge()) {
            LOGGER.warn("tikv Raft too large with region:{}", region);
            return -1;
        } else if (error.hasDiskFull()) {
            LOGGER.warn("tikv DiskFull with region:{}", region);
            return backOffer.nextBackOffTime(BoDiskFull, new TiKVException(error.getMessage()));
        } else if (error.hasRecoveryInProgress()) {
            LOGGER.warn("tikv region recovery in progress:{}", region);
            return backOffer.nextBackOffTime(BoRegionRecoveryInProgress, new TiKVException(error.getMessage()));
        } else if (error.hasMaxTimestampNotSynced()) {
            LOGGER.warn("tikv max timestamp not synced");
            return backOffer.nextBackOffTime(BoMaxTsNotSynced, new TiKVException(error.getMessage()));
        }

        LOGGER.warn("tikv unknown error:{} with region:{}", error, region);
        regionManager.invalidRegion(region);
        regionManager.invalidStore(recv.getStore().getId());
        return -1;
    }

    // ref: https://github.com/tikv/client-go/blob/tidb-5.2/internal/locate/region_request.go#L985
    // OnRegionEpochNotMatch removes the old region and inserts new regions into the cache
    // It returns whether retries the request because it's possible the region epoch is ahead of
    // tikv due to slow appling
    private long onRegionEpochNotMatch(BackOffer backOffer, Errorpb.Error error) {
        TiRegion onRegion = recv.getRegion();
        List<Metapb.Region> currentRegions = error.getEpochNotMatch().getCurrentRegionsList();
        if (currentRegions.isEmpty()) {
            regionManager.invalidRegion(onRegion);
            return -1;
        }

        // find whether the region epoch in recv is ahead of tikv, If so, backoff
        for (Metapb.Region meta : currentRegions) {
            if (meta.getId() == onRegion.getId()
                && (meta.getRegionEpoch().getConfVer() < onRegion.getVerID().getConfVer()
                || meta.getRegionEpoch().getVersion() < onRegion.getVerID().getVer())) {
                String errorMsg = String.format("region epoch:%s is ahead of tikv:%s", onRegion, currentRegions);
                LOGGER.warn(errorMsg);
                return backOffer.nextBackOffTime(BoRegionMiss, new TiKVException(errorMsg));
            }
        }

        boolean needInvalidateOld = true;
        List<TiRegion> newRegions = new ArrayList<>(currentRegions.size());
        // If the region epoch is not ahead of tikv, replace region meta in region cache
        for (Metapb.Region meta : currentRegions) {
            TiRegion newRegion = regionManager.createRegion(meta, null, backOffer);
            newRegions.add(newRegion);
            if (onRegion.getVerID() == newRegion.getVerID()) {
                needInvalidateOld = false;
            }
        }

        if (needInvalidateOld) {
            regionManager.invalidRegion(onRegion);
        }
        for (TiRegion region : newRegions) {
            regionManager.putRegion(region);
        }
        return backOffer.nextBackOffTime(BoTiKVRPC, new RegionException(error));
    }

    @Override
    public boolean handleRequestError(BackOffer backOffer, Exception e) {
        if (recv.onStoreUnreachable(backOffer)) {
            if (!backOffer.canRetryAfterSleep(BoTiKVRPC)) {
                regionManager.invalidRegion(recv.getRegion());
                throw new GrpcException("retry is exhausted", e);
            }
            return true;
        }

        LOGGER.warn("tikv send request failed: {}", e.getMessage());
        if (!backOffer.canRetryAfterSleep(BoTiKVRPC)) {
            regionManager.invalidRegion(recv.getRegion());
            throw new GrpcException(String.format("send tikv request failed: %s", e.getMessage()), e);
        }
        // TiKV maybe down, so do not retry in `callWithRetry`
        // should re-fetch the new leader from PD and send request to it
        return false;
    }

}
